package com.atomgit.ai.proxy.processor;

import com.alibaba.fastjson.JSON;
import com.atomgit.ai.proxy.entity.ModelInfo;
import com.atomgit.ai.proxy.entity.UserApiKey;
import com.atomgit.ai.proxy.param.PromptParam;
import com.atomgit.ai.proxy.param.inner.tongyi.TongYiParam;
import com.atomgit.ai.proxy.param.inner.tongyi.TongYiParamInput;
import com.atomgit.ai.proxy.processor.service.ParamAndResponseProcessor;
import com.atomgit.ai.proxy.vo.ModelAnswerVO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.RequestCallback;
import org.springframework.web.client.RestTemplate;

import javax.servlet.http.HttpServletResponse;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Collections;

/**
 * 通义千问参数与结果处理类
 */
@Slf4j
public class ParamAndResponseProcessorForTongYi implements ParamAndResponseProcessor{
    @Override
    public String getApiKey(ModelInfo modelInfo, UserApiKey userApiKey) {
        return userApiKey.getApiKey();
    }

    @Override
    public String getUrl(ModelInfo modelInfo, UserApiKey userApiKey) {
        return modelInfo.getRequestUrl();
    }

    @Override
    public String generateRequestBody(PromptParam param) {
        TongYiParam tongYiParam = new TongYiParam();
        tongYiParam.setModel(param.getVersion());
        TongYiParamInput tongYiParamInput = new TongYiParamInput();
        BeanUtils.copyProperties(param,tongYiParamInput);
        tongYiParam.setInput(tongYiParamInput);
        return JSON.toJSONString(tongYiParam);
    }

    @Override
    public ModelAnswerVO generateModelAnswerVO(String responseBody) {
        if (responseBody != null) {
            return JSON.parseObject(responseBody, ModelAnswerVO.class);
        }
        return null;
    }

    public String sendRequestInStream(String url, HttpMethod httpMethod, String requestBody, String apiKey, HttpServletResponse response, ParamAndResponseProcessor paramAndResponseProcessor) {
        log.info("sendRequestInStream url:{},httpMethod:{},requestBody:{}",url,httpMethod,requestBody);
        try {
            RestTemplate restTemplate = new RestTemplate();
            // 设置请求头
            MultiValueMap<String, String> headersParam = new HttpHeaders();
            HttpHeaders headers = new HttpHeaders(headersParam);
            headers.setAccept(Collections.singletonList(MediaType.APPLICATION_JSON));
            headers.setContentType(MediaType.APPLICATION_JSON);
            // 添加其他需要的头部信息
            headers.add("Authorization", apiKey);
            headers.add("Accept", "text/event-stream");
            RequestCallback requestCallback = request -> {
                // 设置请求体
                request.getHeaders().addAll(headers);
                request.getBody().write(requestBody.getBytes(StandardCharsets.UTF_8));
            };
            response.setCharacterEncoding("UTF-8");
            restTemplate.execute(URI.create(url), httpMethod, requestCallback, clientHttpResponse -> {
                try (BufferedReader reader = new BufferedReader(new InputStreamReader(clientHttpResponse.getBody(), StandardCharsets.UTF_8))) {
                    // 检车http返回码是否为401(apikey/token失效)
                    if (clientHttpResponse.getStatusCode().value() == 401) {
                        log.info("sendRequest response:{}",response);
                        return "Invalid API-key provided.";
                    }
                    String line;
                    String newDataValue = null;
                    while ((line = reader.readLine()) != null) {
                        char[] buffer = new char[1024 * 5];
                        reader.read(buffer);
                        String result = new String(buffer);
                        int dataIndexFrom = result.indexOf("data:") + 5;
                        String data = result.substring(dataIndexFrom);
                        ModelAnswerVO modelAnswerVO = paramAndResponseProcessor.generateModelAnswerVO(data);
                        response.getWriter().println(JSON.toJSONString(modelAnswerVO));
                        response.getWriter().flush();
                    }
                    return newDataValue;
                }
            });
        } catch (Exception e) {
            log.info(e.getMessage());
            e.printStackTrace();
        }
        return null;
    }


}